#!/usr/bin/env python3
"""
智能图片解压脚本 - 与压缩脚本配套使用
功能特性：
1. 批量解压多个压缩包
2. 智能目录结构恢复
3. 断点续传和跳过已解压
4. 实时进度监控
5. 文件完整性验证
6. 选择性解压（按年份、按文件名）
7. 并行解压优化
"""

import os
import sys
import zipfile
import threading
import time
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Dict, Set, Optional
import argparse
from dataclasses import dataclass
import json
import fnmatch
import hashlib

@dataclass
class DecompressionStats:
    """解压统计信息"""
    total_archives: int = 0
    processed_archives: int = 0
    total_files: int = 0
    extracted_files: int = 0
    total_size_mb: float = 0.0
    extracted_size_mb: float = 0.0
    elapsed_time: float = 0.0
    current_archive: str = ""
    
    def to_dict(self) -> dict:
        return {
            'total_archives': self.total_archives,
            'processed_archives': self.processed_archives,
            'total_files': self.total_files,
            'extracted_files': self.extracted_files,
            'total_size_mb': round(self.total_size_mb, 2),
            'extracted_size_mb': round(self.extracted_size_mb, 2),
            'elapsed_time': round(self.elapsed_time, 2),
            'current_archive': self.current_archive
        }

class SmartImageDecompressor:
    """智能图片解压器"""
    
    def __init__(self, archive_dir: str, output_dir: str, overwrite: bool = False):
        self.archive_dir = Path(archive_dir)
        self.output_dir = Path(output_dir)
        self.overwrite = overwrite
        self.stats = DecompressionStats()
        self.lock = threading.Lock()
        self.start_time = time.time()
        
        # 创建输出目录
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        # 状态文件，用于断点续传
        self.state_file = self.output_dir / "decompression_state.json"
        self.completed_archives = self.load_state()
        
    def load_state(self) -> set:
        """加载已完成的解压状态"""
        if self.state_file.exists():
            try:
                with open(self.state_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    return set(data.get('completed_archives', []))
            except:
                pass
        return set()
    
    def save_state(self, completed_archive: str):
        """保存解压状态"""
        self.completed_archives.add(completed_archive)
        state_data = {
            'completed_archives': list(self.completed_archives),
            'last_update': time.strftime('%Y-%m-%d %H:%M:%S'),
            'stats': self.stats.to_dict()
        }
        with open(self.state_file, 'w', encoding='utf-8') as f:
            json.dump(state_data, f, indent=2, ensure_ascii=False)
    
    def get_archive_files(self, pattern: str = "*.zip") -> List[Path]:
        """获取压缩包文件列表"""
        archives = []
        for archive_file in self.archive_dir.glob(pattern):
            if archive_file.is_file() and archive_file.name not in self.completed_archives:
                archives.append(archive_file)
        return sorted(archives)
    
    def analyze_archive(self, archive_path: Path) -> Dict:
        """分析压缩包内容"""
        try:
            with zipfile.ZipFile(archive_path, 'r') as zipf:
                file_list = zipf.namelist()
                total_size = sum(info.file_size for info in zipf.infolist() if not info.is_dir())
                
                return {
                    'file_count': len([f for f in file_list if not f.endswith('/')]),
                    'total_size': total_size,
                    'file_list': file_list,
                    'valid': True
                }
        except Exception as e:
            return {
                'file_count': 0,
                'total_size': 0,
                'file_list': [],
                'valid': False,
                'error': str(e)
            }
    
    def extract_archive(self, archive_path: Path, year_filter: Optional[List[str]] = None, 
                       file_pattern: Optional[str] = None) -> bool:
        """解压单个压缩包"""
        archive_name = archive_path.name
        
        try:
            print(f"📦 开始解压: {archive_name}")
            
            # 分析压缩包
            analysis = self.analyze_archive(archive_path)
            if not analysis['valid']:
                print(f"❌ 压缩包损坏: {archive_name} - {analysis.get('error', '未知错误')}")
                return False
            
            extracted_count = 0
            skipped_count = 0
            
            with self.lock:
                self.stats.current_archive = archive_name
            
            with zipfile.ZipFile(archive_path, 'r') as zipf:
                for file_info in zipf.infolist():
                    if file_info.is_dir():
                        continue
                    
                    file_path = file_info.filename
                    
                    # 年份过滤
                    if year_filter:
                        year_match = False
                        for year in year_filter:
                            if f"/{year}/" in file_path or file_path.startswith(f"{year}/"):
                                year_match = True
                                break
                        if not year_match:
                            continue
                    
                    # 文件名模式过滤
                    if file_pattern and not fnmatch.fnmatch(os.path.basename(file_path), file_pattern):
                        continue
                    
                    # 目标文件路径
                    target_path = self.output_dir / file_path
                    
                    # 跳过已存在的文件（除非强制覆盖）
                    if target_path.exists() and not self.overwrite:
                        skipped_count += 1
                        continue
                    
                    # 创建目录
                    target_path.parent.mkdir(parents=True, exist_ok=True)
                    
                    # 解压文件
                    with zipf.open(file_info) as source, open(target_path, 'wb') as target:
                        target.write(source.read())
                    
                    extracted_count += 1
                    
                    # 更新统计
                    with self.lock:
                        self.stats.extracted_files += 1
                        self.stats.extracted_size_mb += file_info.file_size / (1024 * 1024)
                    
                    # 每100个文件显示进度
                    if extracted_count % 100 == 0:
                        print(f"     ✅ 已解压 {extracted_count} 个文件...")
            
            print(f"✅ {archive_name} 解压完成: 解压{extracted_count}个文件, 跳过{skipped_count}个文件")
            
            # 保存完成状态
            self.save_state(archive_name)
            
            with self.lock:
                self.stats.processed_archives += 1
            
            return True
            
        except Exception as e:
            print(f"❌ {archive_name} 解压失败: {e}")
            return False
    
    def verify_extraction(self, archive_path: Path) -> Dict:
        """验证解压完整性"""
        try:
            verification_result = {
                'archive': archive_path.name,
                'total_files': 0,
                'verified_files': 0,
                'missing_files': [],
                'size_mismatches': [],
                'success': True
            }
            
            with zipfile.ZipFile(archive_path, 'r') as zipf:
                for file_info in zipf.infolist():
                    if file_info.is_dir():
                        continue
                    
                    verification_result['total_files'] += 1
                    target_path = self.output_dir / file_info.filename
                    
                    if not target_path.exists():
                        verification_result['missing_files'].append(file_info.filename)
                        verification_result['success'] = False
                    else:
                        # 检查文件大小
                        actual_size = target_path.stat().st_size
                        expected_size = file_info.file_size
                        
                        if actual_size != expected_size:
                            verification_result['size_mismatches'].append({
                                'file': file_info.filename,
                                'expected': expected_size,
                                'actual': actual_size
                            })
                            verification_result['success'] = False
                        else:
                            verification_result['verified_files'] += 1
            
            return verification_result
            
        except Exception as e:
            return {
                'archive': archive_path.name,
                'success': False,
                'error': str(e)
            }
    
    def print_progress(self):
        """打印实时进度"""
        while True:
            time.sleep(3)  # 每3秒更新一次
            
            with self.lock:
                elapsed = time.time() - self.start_time
                self.stats.elapsed_time = elapsed
                
                # 计算速度
                files_per_sec = self.stats.extracted_files / elapsed if elapsed > 0 else 0
                mb_per_sec = self.stats.extracted_size_mb / elapsed if elapsed > 0 else 0
                archives_per_min = self.stats.processed_archives / (elapsed / 60) if elapsed > 0 else 0
                
                print(f"\n📊 解压进度:")
                print(f"   压缩包: {self.stats.processed_archives}/{self.stats.total_archives} "
                      f"({self.stats.processed_archives/self.stats.total_archives*100:.1f}%)" if self.stats.total_archives > 0 else "   压缩包: 统计中...")
                print(f"   文件: {self.stats.extracted_files}/{self.stats.total_files} "
                      f"({self.stats.extracted_files/self.stats.total_files*100:.1f}%)" if self.stats.total_files > 0 else "   文件: 统计中...")
                print(f"   大小: {self.stats.extracted_size_mb:.1f}MB/{self.stats.total_size_mb:.1f}MB")
                print(f"   速度: {files_per_sec:.1f}文件/秒, {mb_per_sec:.1f}MB/秒, {archives_per_min:.1f}包/分钟")
                print(f"   耗时: {elapsed/60:.1f}分钟")
                print(f"   当前: {self.stats.current_archive}")
                
            if self.stats.processed_archives >= self.stats.total_archives and self.stats.total_archives > 0:
                break
    
    def list_archives(self, pattern: str = "*.zip"):
        """列出压缩包信息"""
        archives = self.get_archive_files(pattern)
        
        print(f"📋 发现 {len(archives)} 个压缩包:")
        print("-" * 80)
        
        total_files = 0
        total_size_mb = 0
        
        for i, archive_path in enumerate(archives, 1):
            analysis = self.analyze_archive(archive_path)
            size_mb = analysis['total_size'] / (1024 * 1024)
            status = "✅ 有效" if analysis['valid'] else "❌ 损坏"
            completed = "✅ 已完成" if archive_path.name in self.completed_archives else "⏳ 待处理"
            
            print(f"{i:2d}. {archive_path.name}")
            print(f"    文件数: {analysis['file_count']}, 大小: {size_mb:.1f}MB, 状态: {status}, {completed}")
            
            if analysis['valid']:
                total_files += analysis['file_count']
                total_size_mb += size_mb
        
        print("-" * 80)
        print(f"📊 总计: {total_files} 个文件, {total_size_mb:.1f}MB")
    
    def run(self, max_workers: int = 2, year_filter: Optional[List[str]] = None, 
           file_pattern: Optional[str] = None, verify: bool = False, 
           list_only: bool = False, pattern: str = "*.zip"):
        """运行解压任务"""
        print("🚀 智能图片解压器启动")
        print(f"📂 压缩包目录: {self.archive_dir}")
        print(f"📁 输出目录: {self.output_dir}")
        print(f"🔄 覆盖模式: {'是' if self.overwrite else '否'}")
        
        if year_filter:
            print(f"📅 年份过滤: {', '.join(year_filter)}")
        if file_pattern:
            print(f"🔍 文件模式: {file_pattern}")
        
        # 获取压缩包列表
        archives = self.get_archive_files(pattern)
        
        if not archives:
            print("❌ 未找到待解压的压缩包")
            return
        
        # 仅列出信息
        if list_only:
            self.list_archives(pattern)
            return
        
        # 分析所有压缩包
        print(f"\n📋 分析 {len(archives)} 个压缩包...")
        self.stats.total_archives = len(archives)
        
        for archive_path in archives:
            analysis = self.analyze_archive(archive_path)
            if analysis['valid']:
                self.stats.total_files += analysis['file_count']
                self.stats.total_size_mb += analysis['total_size'] / (1024 * 1024)
        
        print(f"📊 总计: {self.stats.total_files} 个文件, {self.stats.total_size_mb:.1f}MB")
        
        # 启动进度监控线程
        progress_thread = threading.Thread(target=self.print_progress, daemon=True)
        progress_thread.start()
        
        # 并行解压
        print(f"\n📦 开始并行解压 (最多{max_workers}个压缩包同时进行)...")
        
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # 提交解压任务
            future_to_archive = {
                executor.submit(self.extract_archive, archive, year_filter, file_pattern): archive 
                for archive in archives
            }
            
            # 等待完成
            success_count = 0
            for future in as_completed(future_to_archive):
                archive = future_to_archive[future]
                try:
                    success = future.result()
                    if success:
                        success_count += 1
                        print(f"🎉 {archive.name} 解压成功")
                    else:
                        print(f"💥 {archive.name} 解压失败")
                except Exception as e:
                    print(f"💥 {archive.name} 解压异常: {e}")
        
        # 验证解压结果
        if verify:
            print(f"\n🔍 验证解压完整性...")
            verification_results = []
            
            for archive_path in archives:
                if archive_path.name in self.completed_archives:
                    result = self.verify_extraction(archive_path)
                    verification_results.append(result)
                    
                    if result['success']:
                        print(f"✅ {result['archive']}: 验证通过 ({result['verified_files']}/{result['total_files']})")
                    else:
                        print(f"❌ {result['archive']}: 验证失败")
                        if result.get('missing_files'):
                            print(f"   缺失文件: {len(result['missing_files'])} 个")
                        if result.get('size_mismatches'):
                            print(f"   大小不匹配: {len(result['size_mismatches'])} 个")
        
        # 最终统计
        total_time = time.time() - self.start_time
        print(f"\n🏁 解压完成!")
        print(f"📊 最终统计:")
        print(f"   成功解压: {success_count}/{len(archives)} 个压缩包")
        print(f"   解压文件: {self.stats.extracted_files} 个")
        print(f"   解压大小: {self.stats.extracted_size_mb:.1f}MB")
        print(f"   总耗时: {total_time/60:.1f}分钟")
        print(f"   平均速度: {self.stats.extracted_files/(total_time/60):.0f}文件/分钟")

def main():
    parser = argparse.ArgumentParser(description="智能图片解压器")
    parser.add_argument("archive_dir", help="压缩包目录")
    parser.add_argument("-o", "--output", default="./extracted", help="输出目录 (默认: ./extracted)")
    parser.add_argument("-w", "--workers", type=int, default=2, help="并行线程数 (默认: 2)")
    parser.add_argument("--overwrite", action="store_true", help="覆盖已存在的文件")
    parser.add_argument("--verify", action="store_true", help="验证解压完整性")
    parser.add_argument("--list", action="store_true", help="仅列出压缩包信息")
    parser.add_argument("--years", nargs="+", help="仅解压指定年份的文件")
    parser.add_argument("--pattern", default="*.zip", help="压缩包文件模式 (默认: *.zip)")
    parser.add_argument("--file-pattern", help="文件名过滤模式 (如: *.jpg)")
    
    args = parser.parse_args()
    
    if not os.path.exists(args.archive_dir):
        print(f"❌ 压缩包目录不存在: {args.archive_dir}")
        sys.exit(1)
    
    decompressor = SmartImageDecompressor(
        archive_dir=args.archive_dir,
        output_dir=args.output,
        overwrite=args.overwrite
    )
    
    try:
        decompressor.run(
            max_workers=args.workers,
            year_filter=args.years,
            file_pattern=args.file_pattern,
            verify=args.verify,
            list_only=args.list,
            pattern=args.pattern
        )
    except KeyboardInterrupt:
        print("\n⏹️  用户中断，状态已保存，下次运行将从断点继续")
    except Exception as e:
        print(f"\n💥 解压过程出错: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()
